package com.tutucms.common.tree;

import org.apache.commons.collections.CollectionUtils;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 树构造器
 * Created by wangxudong on 2020/3/28.
 */
public class TreeBuilder {
    private TreeBuilder() {

    }

    /**
     * 根据列表构造树，这里的树可以是多个根的。
     *
     * @param treeNodeList 列表
     * @param sort         是否排序
     * @return 构造后的树
     */
    public static List<? extends TreeNode> build(List<? extends TreeNode> treeNodeList, boolean sort) {
        List<TreeNode> trees = new ArrayList<>();
        Map<String, TreeNode> treeNodeMap = new HashMap<>(treeNodeList.size() * 2);
        for (TreeNode treeNode : treeNodeList) {
            treeNodeMap.put(treeNode.getId(), treeNode);
        }
        for (Map.Entry entry : treeNodeMap.entrySet()) {
            TreeNode treeNode = (TreeNode) entry.getValue();
            TreeNode parentNode = treeNodeMap.get(treeNode.getParentId());
            if (parentNode == null) {
                trees.add(treeNode);
            } else {
                if (parentNode.getChildren() == null) {
                    parentNode.setChildren(new ArrayList<>());
                }
                parentNode.getChildren().add(treeNode);
            }
        }

        if (sort) {
            treeNodeList.stream().filter(treeNode -> !CollectionUtils.isEmpty(treeNode.getChildren()))
                    .forEach(treeNode -> Collections.sort(treeNode.getChildren()));
            Collections.sort(trees);
        }
        return trees;
    }

    /**
     * 在树节点列表中查找子树
     *
     * @param treeNodeList 树节点列表
     * @param id           待查找节点id
     * @return 子树
     */
    public static TreeNode getSubTreeFromList(List<? extends TreeNode> treeNodeList, String id) {
        TreeNode targetNode = null;
        Map<String, TreeNode> treeNodeMap = new HashMap<>(treeNodeList.size() * 2);
        for (TreeNode treeNode : treeNodeList) {
            treeNodeMap.put(treeNode.getId(), treeNode);
        }
        for (Map.Entry entry : treeNodeMap.entrySet()) {
            TreeNode treeNode = (TreeNode) entry.getValue();
            if (treeNode.getId().equals(id)) {
                targetNode = treeNode;
            }
            TreeNode parentNode = treeNodeMap.get(treeNode.getParentId());
            if (parentNode != null) {
                if (parentNode.getChildren() == null) {
                    parentNode.setChildren(new ArrayList<>());
                }
                parentNode.getChildren().add(treeNode);
            }
        }
        return targetNode;
    }

    /**
     * 将树形转成列表
     *
     * @param root 树节点
     * @return 树节点列表
     */
    public static List<? extends TreeNode> treeToList(TreeNode root) {
        if (root == null) {
            return null;
        }
        List<TreeNode> treeNodeList = new ArrayList<>();
        treeNodeList.add(root);
        List<TreeNode> children = root.getChildren();
        if (!CollectionUtils.isEmpty(children)) {
            for (TreeNode child : children) {
                treeNodeList.addAll(treeToList(child));
            }
        }
        return treeNodeList;
    }

    public static List<? extends TreeNode> treeToList(List<? extends TreeNode> trees) {
        if (trees == null) {

            return null;
        }
        List<TreeNode> treeNodeList = new ArrayList<>();
        for (TreeNode node : trees) {
            treeNodeList.add(node);
            if (node.getChildren() != null) {
                treeNodeList.addAll(treeToList(node.getChildren()));
                node.setChildren(null);
            }
        }

        return treeNodeList;
    }

    /**
     * 获取所有的祖先
     *
     * @param treeNodeList
     * @param ids
     * @return
     */
    public static List<? extends TreeNode> getAllAncestors(List<? extends TreeNode> treeNodeList, List<String> ids) {
        Map<String, TreeNode> treeNodeMap = new HashMap<>(treeNodeList.size() * 2);
        treeNodeList.stream().forEach(treeNode -> treeNodeMap.put(treeNode.getId(), treeNode));
        List<TreeNode> ancestors = new ArrayList<>();
        ids.stream().forEach(id -> {
            TreeNode node = treeNodeMap.get(id);
            if (node == null) {
                throw new RuntimeException("cannot find this node");
            }
            while (true) {
                node = treeNodeMap.get(node.getParentId());
                if (node == null) {
                    break;
                } else {
                    ancestors.add(node);
                }
            }
        });

        return ancestors.stream().distinct().collect(Collectors.toList());
    }

    /**
     * 获取所有的后代
     *
     * @param treeNodeList
     * @param ids
     * @return
     */
    public static List<? extends TreeNode> getAllDescendants(List<? extends TreeNode> treeNodeList, List<String> ids) {
        Map<String, List<TreeNode>> descendantsMap = new HashMap<>(treeNodeList.size());
        Map<String, TreeNode> treeNodeMap = new HashMap<>(treeNodeList.size());
        treeNodeList.stream().forEach(treeNode -> {
            List<TreeNode> descendants = descendantsMap.get(treeNode.getParentId());
            if (descendants == null) {
                descendants = new ArrayList<>();
                descendantsMap.put(treeNode.getParentId(), descendants);
            }
            descendants.add(treeNode);
            treeNodeMap.put(treeNode.getId(), treeNode);
        });

        List<TreeNode> allDescendants = new ArrayList<>();
        List<TreeNode> currentDescendants = ids.stream().map(treeNodeMap::get).collect(Collectors.toList());
        List<TreeNode> toFindDescendants = new ArrayList<>();
        while (!currentDescendants.isEmpty()) {
            for (TreeNode currentDescendant : currentDescendants) {
                List<TreeNode> oneToFind = descendantsMap.get(currentDescendant.getId());
                if (!CollectionUtils.isEmpty(oneToFind)) {
                    toFindDescendants.addAll(oneToFind);
                    allDescendants.addAll(oneToFind);
                }
            }
            currentDescendants = toFindDescendants;
            toFindDescendants = new ArrayList<>();
        }
        return allDescendants;
    }

}
